package com.example.rpc.demo.discovery;

import com.example.rpc.demo.consts.RpcConstant;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooKeeper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ThreadLocalRandom;

public class ServiceDiscovery {

    private static final Logger logger = LoggerFactory.getLogger(ServiceDiscovery.class);

    private CountDownLatch latch = new CountDownLatch(1);

    private volatile Map<String, List<String>> serviceMap = new HashMap<>();

    private String registryAddress;

    public ServiceDiscovery(String registryAddress) {
        this.registryAddress = registryAddress;
        ZooKeeper zk  = connectServer();
        if(zk!=null) {
            watchNode(zk);
        }

    }

    private ZooKeeper connectServer() {
        ZooKeeper zk = null;
        try {
            zk = new ZooKeeper(registryAddress,RpcConstant.ZK_SESSION_TIMEOUT,event -> {
                if(event.getState() == Watcher.Event.KeeperState.SyncConnected) {
                    latch.countDown();
                }
            });

            latch.await();
        } catch (InterruptedException | IOException e) {
            logger.error("connect zk server failed." ,e);
        }
        return zk;
    }

    private void watchNode(ZooKeeper zk) {


        try {
            List<String> serviceList =  zk.getChildren(RpcConstant.ZK_REGISTRY_PATH,event -> {
                if(event.getType() == Watcher.Event.EventType.NodeChildrenChanged) {
                    watchNode(zk);
                }
            });

            Map<String, List<String>> serviceMap = new HashMap<>();
            for(String service: serviceList) {
                List<String> nodeList = zk.getChildren(RpcConstant.ZK_REGISTRY_PATH+"/"+service,event -> {
                   if(event.getType() == Watcher.Event.EventType.NodeChildrenChanged) {
                       watchNode(zk);
                   }
                });

                List<String> addressList = new ArrayList<>();
                for(String node:nodeList) {
                    byte[] bytes = zk.getData(RpcConstant.ZK_REGISTRY_PATH + "/" + service + "/" + node, false, null);
                    addressList.add(new String(bytes));
                }

                serviceMap.put(service, addressList);
            }

            logger.info("node data : {}",serviceMap);

            this.serviceMap = serviceMap;
        } catch (KeeperException | InterruptedException e) {
            logger.error("zk watch node failed . ",e);
        }

    }

    public String discover(String serviceName) {
        String data = null;
        if(serviceMap.containsKey(serviceName)) {
            List<String> addressList = serviceMap.get(serviceName);
            int size = addressList.size();
            if(size == 1) {
                data = addressList.get(0);
            } else {
                data = addressList.get(ThreadLocalRandom.current().nextInt(size));
            }
        }

        return data;
    }





}
